#include <cmath>

#include "FEMModel.h"
#include "TetrahedronState.h"
#include "Setting.h"

/*
Viscosity along the edges of a tetrahedron.

TODO: flop count
*/
void SimpleSpringViscosityForce (TetrahedronState const*me,
							   Real *force_vec,
							   Real const* velocity,
							   Real const* displacement_vec,
							   Real viscosity,
							   FiniteElementModel*fem)
{
	Vector3 x[4];
	Vector3 v[4];

	Real const * locs = fem->node_locations_.AccessArray ();
	for (int i = 4; i--;)
	{
		for (int j  =  3; j--;)
		{
			int k =3*me->incident_nodes_[i]+ j;
			x[i](j) = locs[k] + displacement_vec[k];
			v[i](j) = velocity[k];
		}
	}

	for (int i=4; i--;)
		for (int j = 0 ; j < i; j++)
		{
			Vector3 diff = x[i]  - x[j];
			diff /= diff.Length ();

			/*
			Scale with volume to be independent of edge-tet incidence.
			*/
			Real f = (diff * (v[i] - v[j]))* viscosity *me->volume_;

			for (int k = 3; k--; )
			{
				force_vec[3*me->incident_nodes_[i]+ k] += -f * diff (k);
				force_vec[3*me->incident_nodes_[j]+ k] +=  f * diff (k)  ;
			}
		}
}

/*
Elastic forces assuming that each edge has a linear spring.

TODO: add spring constant as a parameter.

COST:

(+ 12 (* 6 (+ 3 5 3 7 7 3))) = 180 flops
*/
int SimpleSpringForce (TetrahedronState const *me,
					 Real *force_vec,
					 Real const*  displacement_vec,
					 FiniteElementModel*fem)
{
	Real young =  GetNumberSetting ("young");
	Vector3 X[4];
	Vector3 x[4];

	Real const * locs = fem->node_locations_.AccessArray ();
	for (int i = 4; i--;)
	{
		for (int j  =  3; j--;)
		{
			int k =3*me->incident_nodes_[i]+ j;
			X[i](j) = locs[k];
			x[i](j) = locs[k] + displacement_vec[k];
		}
	}

	for (int i=4; i--;)
		for (int j = 0 ; j < i; j++)
		{
			Real L2 = (X[i] - X[j]).Length2 ();

			Vector3 diff = x[i]  - x[j];
			Real l = diff.Length ();

			Real f = young *  (me->volume_ / L2) * (l - sqrt (L2)) / l  ;
			for (int k = 3; k--; )
			{
				force_vec[3*me->incident_nodes_[i]+ k] += -f * diff (k);
				force_vec[3*me->incident_nodes_[j]+ k] +=  f * diff (k)  ;
			}
		}

		return  180;
}
